package pa14.asmt.asmt3.main;

import java.util.HashSet;
import java.util.List;
import java.util.Set;

import soot.G;
import soot.SootMethod;
import soot.Unit;
import soot.Value;
import soot.ValueBox;
import soot.jimple.AssignStmt;
import soot.jimple.InvokeExpr;
import soot.jimple.InvokeStmt;
import soot.toolkits.graph.ExceptionalUnitGraph;
import soot.toolkits.scalar.ForwardFlowAnalysis;

public class TaintVariableAnalysis extends ForwardFlowAnalysis<Unit, Set<Value>>{
	private Set<Value> taintedVars;
	private Set<Value> taintedVarsLiveAtSink;
	private Set<Value> entryTaintedVars;
	private boolean debuglevel;
	
	public Set<Value> getTaintedVars(){
		return (null != taintedVars)?taintedVars: new HashSet<Value>();
	}
	
	
	public TaintVariableAnalysis(ExceptionalUnitGraph graph, Set<Value> t, boolean debuglevel) {
		super(graph);
		this.entryTaintedVars = t;
		this.taintedVars = new HashSet<Value>();
		this.taintedVarsLiveAtSink = new HashSet<Value>();
		this.debuglevel = debuglevel;
		//	Calculate gen and kill sets prior
		
		
		//	start analysis
		doAnalysis();
	}

	@Override
	protected void flowThrough(Set<Value> in, Unit thisStmt, Set<Value> out) {
		// TODO Flow through method
		if(debuglevel)
			G.v().out.println("The ins right now:" + in);
		
		out.removeAll(out);	//remove everything
		
		if(isTaintSink(thisStmt)){
			//	See if there are any variables that are used in the print statement
			//	and are also in the in set.
			
			
			//get the used variables
			List<ValueBox> usedVars = thisStmt.getUseBoxes();
			for(ValueBox var: usedVars){
				Value v = var.getValue();
				if(in.contains(v)){
					taintedVars.add(v);
				}
			}
			
			taintedVarsLiveAtSink.addAll(removeJimpleTemps(in));
		}
		
		//List<ValueBox> uses = thisStmt.getUseBoxes();
		//List<ValueBox> defs = thisStmt.getDefBoxes();
		Set<Value> gen = new HashSet<Value>();
		Set<Value> kill = new HashSet<Value>();
		
		// if any of the right ops are already tainted, then the left op is put in the gen
		if (thisStmt instanceof AssignStmt){
			if(debuglevel){
				G.v().out.println("\nPrinting lefts and rights of opboxes only for assignment statements!");
				G.v().out.println("\n\t\tStatement is: " + thisStmt);
			}
			AssignStmt stm = (AssignStmt)thisStmt;
			List<ValueBox> rights = thisStmt.getUseBoxes();
			Value left = stm.getLeftOp();

			for(ValueBox vb: rights){
				if(in.contains(vb.getValue())){
					gen.add(left);
					break;
				}
			}
			if(debuglevel)
				G.v().out.println("\n Gen set now:" + gen);
			
			//	calculate the kill sets
			//	variable will only be killed when it gets a new definition,
			//	and one which does not rely on any tainted variables
			boolean hasNoTaintedVars = true;
			for(ValueBox vb: rights){
				if(in.contains(vb.getValue())){
					hasNoTaintedVars = false;
					break;
				}
			}
			if(hasNoTaintedVars){
				kill.add(left);
			}
			
		}
		
		out.addAll(in);
		out.removeAll(kill);
		out.addAll(gen);
	}

	@Override
	protected void copy(Set<Value> source, Set<Value> dest) {
		dest.removeAll(dest);
		dest.addAll(source);
	}

	@Override
	protected void merge(Set<Value> in1, Set<Value> in2, Set<Value> out) {
		out.removeAll(out);
		out.addAll(in1);
		out.addAll(in2);		//	the merge operator is union
								//	because it is a may analysis
	}

	
	@Override
	protected Set<Value> entryInitialFlow() {
		// entry values
		// should be equal to the values of the parameters.
		// All parameters are by default, tainted.
		
		//	return the parameters that we have precalculated
		return this.entryTaintedVars;
	}

	@Override
	protected Set<Value> newInitialFlow() {
		// initial values for all boxes
		// should be null
		return new HashSet<Value>();
	}
	
	private boolean isTaintSink(Unit u){
		if(u instanceof InvokeStmt) {
		    InvokeExpr e = ((InvokeStmt) u).getInvokeExpr();
		    SootMethod m = e.getMethod();
		    if(m.getName().equals("println") && m.getDeclaringClass().getName().equals("java.io.PrintStream"))
		    	return true;
		}

		return false;
	}


	public Set<Value> getTaintedVarsLiveAtSink() {
		return taintedVarsLiveAtSink;
	}

	public Set<Value> removeJimpleTemps(Set<Value> in){
		Set<Value> ReducedSet = new HashSet<Value>();
		
		for(Value v : in){
			String strV = v.toString();
			if(strV.length()>5){
				if(!strV.substring(0, 5).equals("temp$"))
					ReducedSet.add(v);	
			}else{
				ReducedSet.add(v);
			}
		}
		return ReducedSet;
	}
}
